import dgl
import torch
import numpy as np
import pandas as pd
from Config import Config

"""
对数据集进行预处理，将mRNA分子整合成Graph的形式
可以设置为成对的碱基加边，设置是否使用密码子
"""


class GraphParser(object):
    def __init__(self, config, mode="train"):
        self.config = config
        self.mode = mode

    def parse(self, sample):
        labels = []  # 初始化节点
        nfeat = []  # 初始化节点特征
        efeat = []  # 初始化边特征
        edges = []  # 初始化边
        loop_type = []  # 初始化结构信息
        train_mask = []  # 初始化训练mask
        test_mask = []  # 初始化测试mask

        sequence = sample['sequence']  # RNA序列，由A、U、G、C组成
        predicted_loop_type = sample['predicted_loop_type']  # 循环类型，可以理解为RNA分子的结构特征
        seq_length = sample['seq_length']  # RNA序列的长度
        seq_scored = sample['seq_scored']  # 需要预测的RNA序列长度

        pair_info = match_pair(sample['structure'])  # 描述一个碱基是成对还是不成对的，即结构信息

        paired_nodes = {}
        for j in range(seq_length):
            add_base_node(nfeat, sequence[j])
            add_loop_type(loop_type, predicted_loop_type[j])

            if j + 1 < seq_length:  # edge between current node and next node
                # 判断当前节点和下一个节点是否有边
                add_edges_between_base_nodes(edges, efeat, j, j + 1)

            if pair_info[j] != -1:
                # 若根据structure，两个节点成对，则将节点保存到paired_nodes中
                # {0:[5, 13], ...} 意思是 5号节点 和 13号节点 成对
                if pair_info[j] not in paired_nodes:
                    paired_nodes[pair_info[j]] = [j]
                else:
                    paired_nodes[pair_info[j]].append(j)

            # 若当前遍历节点小于需要预测的长度，则将当前train_mask置为True
            train_mask.append(j < seq_scored)
            test_mask.append(True)

        # 若为训练模式，则为预测的序列加入数据
        if self.mode != 'test':
            react = sample['reactivity']
            deg_Mg_pH10 = sample['deg_Mg_pH10']
            deg_Mg_50C = sample['deg_Mg_50C']

            for j in range(seq_length):
                if j < seq_scored:  # 若是指定需要预测的地方，加入数据
                    labels.append([react[j], deg_Mg_pH10[j], deg_Mg_50C[j]])
                else:  # 否则填充0
                    labels.append([0, 0, 0])

        # 若配置文件中指定，需要为成对的节点加边
        if self.config.add_edge_for_paired_nodes:
            for pair in paired_nodes.values():
                add_edges_between_paired_nodes(edges, efeat, pair[0], pair[1])

        # 若配置文件中指定，需要加入密码子
        if self.config.add_codon_nodes:
            codon_node_idx = seq_length - 1
            for j in range(seq_length):
                if j % 3 == 0:
                    # add codon node
                    # 添加密码子
                    add_codon_node(nfeat)
                    add_loop_type(loop_type, predicted_loop_type='codon')
                    codon_node_idx += 1  # 密码子下标加1
                    train_mask.append(False)
                    test_mask.append(False)
                    if self.mode != "test":
                        labels.append([0, 0, 0])

                    if codon_node_idx > seq_length:  # 保证codon_node_idx - 1密码子节点存在
                        # add edges between adjacent codon nodes
                        # 在两个密码子节点之间加边
                        add_edges_between_codon_nodes(edges, efeat, codon_node_idx - 1, codon_node_idx)

                # add edges between codon node and base node
                # 密码子已生成，在密码子和基本节点之间加边
                add_edges_between_codon_and_base_node(edges, efeat, j, codon_node_idx)

        gdata = {}
        gdata['nfeat'] = np.array(nfeat, dtype="int64")
        gdata['loop_type'] = np.array(loop_type, dtype="int64")
        gdata['edges'] = np.array(edges, dtype="int64")
        gdata['efeat'] = np.array(efeat, dtype="int64")
        if self.mode != "test":
            gdata['labels'] = np.array(labels, dtype="float32")
            gdata['mask'] = np.array(train_mask, dtype=bool).reshape(-1, 1)
        else:
            # fake labels
            gdata['labels'] = np.zeros((self.config.batch_size, self.config.num_class))
            gdata['mask'] = np.array(test_mask, dtype=bool).reshape(-1, 1)

        return gdata


def add_base_node(node_features, sequence):
    node_dict = {'A': 0, 'U': 1, 'G': 2, 'C': 3}
    node_features.append(node_dict[sequence])


def add_loop_type(loop_type, predicted_loop_type):
    loop_type_dict = {'S': 0, 'M': 1, 'I': 2, 'B': 3, 'H': 4, 'E': 5, 'X': 6, 'codon': 7}
    loop_type.append(loop_type_dict[predicted_loop_type])


def add_codon_node(node_features):
    feature = 4
    node_features.append(feature)


# add directed edge for node1 -> node2 and for node2 -> node1
# 使用双向图
def add_edges(edge_index, edge_features, node1, node2, feature1, feature2):
    edge_index.append([node1, node2])
    edge_features.append(feature1)
    edge_index.append([node2, node1])
    edge_features.append(feature2)


def add_edges_between_base_nodes(edge_index, edge_features, node1, node2):
    """在指定的两个节点之间添加边"""
    edge_feature1 = edge_feature2 = 0
    add_edges(edge_index, edge_features, node1, node2,
              edge_feature1, edge_feature2)


def add_edges_between_paired_nodes(edge_index, edge_features, node1, node2):
    # 碱基对之间的边
    edge_feature1 = edge_feature2 = 1
    add_edges(edge_index, edge_features, node1, node2,
              edge_feature1, edge_feature2)


def add_edges_between_codon_and_base_node(edge_index, edge_features, node1, node2):
    # 密码子与碱基之间的边
    edge_feature1 = edge_feature2 = 2
    add_edges(edge_index, edge_features, node1, node2,
              edge_feature1, edge_feature2)


def add_edges_between_codon_nodes(edge_index, edge_features, node1, node2):
    # 密码子之间的边
    edge_feature1 = edge_feature2 = 3
    add_edges(edge_index, edge_features, node1, node2,
              edge_feature1, edge_feature2)


def match_pair(structure):
    """将结构信息转为数字信息"""
    """
    若两个碱基无成对连接，则在pair_no_stack中添加-1
    若两个碱基成对连接，则将两个碱基按原位置添加相同的标号
    """
    pair = [-1] * len(structure)
    pair_no = -1

    pair_no_stack = []
    for i, c in enumerate(structure):
        if c == '(':
            pair_no += 1
            pair[i] = pair_no
            pair_no_stack.append(pair_no)
        elif c == ')':
            pair[i] = pair_no_stack.pop()
    return pair


if __name__ == '__main__':
    # 读取数据集, 进行测试
    data_path = "../dataset/valid.json"
    data_set = pd.read_json(data_path, lines=True)
    sample = data_set.loc[0]
    print(sample)

    pair = match_pair(sample['structure'])
    print(pair)

    model_name = 'GCN'
    device = None

    config = Config(model_name, device)
    parser = GraphParser(config=config, mode='train')
    data = []
    gdata = parser.parse(sample)
    print(gdata)
    u = gdata['edges'][:, 0]
    v = gdata['edges'][:, 1]

    u = torch.tensor(u)
    v = torch.tensor(v)
    # 构造图
    graph = dgl.graph((u, v))

    graph.ndata['nfeat'] = gdata['nfeat']
    graph.ndata['mask'] = gdata['mask']
    graph.edata['efeat'] = gdata['efeat']

    print(graph)
